from puzzle import Puzzle

class Menu:
    def __init__(self, menu_options, game_state, parent_menu = None, parent_screen = None, cursor_index = 0):
        # menu_options is a list of dictionaries with keys "option_name" and "on_select"
        self.menu_options = menu_options
        self.game_state = game_state
        self.parent_menu = parent_menu
        self.parent_screen = parent_screen
        self.selected_option = cursor_index

    def navigate(self, direction):
        self.selected_option = (self.selected_option + direction) % len(self.menu_options)

    def cancel(self):
        if self.parent_menu is None and self.parent_screen is None:
            return
        self.game_state.menu = self.parent_menu
        if self.parent_menu is None and self.game_state.screen == 'menu':
            self.game_state.screen = self.parent_screen

    def force_screen(self, screen_name):
        if self.game_state.screen == 'menu' and self.game_state.menu == self:
            self.game_state.menu = None
            self.game_state.screen = screen_name

    def select(self):
        option = self.menu_options[self.selected_option]
        if 'on_select' in option:
            option['on_select'](option, self, self.game_state)

    def open_sub_menu(self, new_options, cursor_index = 0):
        sub_menu = Menu(new_options, self.game_state, self, cursor_index = cursor_index)
        self.game_state.menu = sub_menu

def save_and_exit(option, menu, game_state):
    menu.force_screen('exit_to_os')

def quit_to_main_menu(option, menu, game_state):
    game_state.open_main_menu()

def new_main_menu(game_state_parent):
    def on_select_resume_story(option, menu, game_state):
        game_state.mode = 'story'
        game_state.level = None
        story_info = option['meta_info']
        story = story_info['story']
        current_level = story_info['level']
        level_history = story_info['level_history']
        game_state.load_level_history(level_history)
        if 'current_game' in story_info:
            current_game = story_info['current_game']
            saved_puzzle = Puzzle(current_game['quote'], current_game['key'], current_game['letter_mappings'])
            game_state.load_story_level(story, current_level, True)
            game_state.load_puzzle(saved_puzzle)
        else:
            game_state.load_story_level(story, current_level, False)

    def on_select_new_game_story(option, menu, game_state):
        game_state.mode = 'story'
        game_state.level = None
        story = option['meta_info']
        init_level_id = story['init_level']
        game_state.load_level_history({"unlocked_levels": [], "completed_levels": []})
        game_state.load_story_level(story, init_level_id, False)

    def on_select_level(option, menu, game_state):
        game_state.mode = 'story'
        game_state.level = None
        story_info = option["meta_info"]
        story = story_info["story"]
        level = option["option_value"]
        level_history = story_info["level_history"]
        game_state.load_level_history(level_history)
        game_state.load_story_level(story, level["id"], False)

    def on_open_level_select_sub_menu(option, menu, game_state):
        story_info = option['meta_info']
        story = story_info['story']
        levels = story['levels']
        level_history = story_info['level_history']
        unlocked_levels = level_history['unlocked_levels']

        # Convert the list of levels into a dictionary for quick lookup
        level_dict = { level["id"]: level for level in levels}

        menu_options = []
        for level_id in unlocked_levels:
            level = level_dict[level_id]
            menu_options.append({ "option_name": level['name'], "on_select": on_select_level, "option_value": level, "meta_info": story_info })

        menu_options.sort(key=lambda x: x["option_value"]["index"])
        menu.open_sub_menu(menu_options)

    def on_select_story(option, menu, game_state):
        story = option['option_value']
        saved_game_data = game_state.data_manager.try_load_game()
        if 'story_mode' in saved_game_data and story['id'] in saved_game_data['story_mode']:
            saved_story_progress = saved_game_data['story_mode'][story['id']]
            current_level = saved_story_progress['level']
            level_history = saved_story_progress['level_history']

            story_info = {
                "story": story,
                "saved_story_progress": saved_story_progress,
                "current_level": current_level,
                "level_history": level_history,
                "level": current_level
            }
            if 'current_game' in saved_story_progress:
                story_info['current_game'] = saved_story_progress['current_game']

            menu_options = [
                { "option_name": "Resume Game", "on_select": on_select_resume_story, "meta_info": story_info },
                { "option_name": "Level Select", "on_select": on_open_level_select_sub_menu, "meta_info": story_info }
            ]
        else:
            menu_options = [{ "option_name": "New Game", "on_select": on_select_new_game_story, "meta_info": story}]
        menu.open_sub_menu(menu_options)

    def story_mode(option, menu, game_state):
        story_list = [story for story in game_state.data_manager.list_stories()]
        story_options = [{ "option_name": story['name'], "option_value": story, "on_select": on_select_story } for story in story_list]
        menu.open_sub_menu(story_options)

    def endless_mode(option, menu, game_state):
        game_state.mode = 'endless'
        saved_game_data = game_state.data_manager.try_load_game()
        if 'game' in saved_game_data:
            saved_puzzle = Puzzle(saved_game_data['game']['quote'], saved_game_data['game']['key'], saved_game_data['game']['letter_mappings'])
            game_state.load_puzzle(saved_puzzle)
            game_state.screen = 'puzzle'
        else:
            game_state.start_new_puzzle()
            game_state.screen = 'puzzle'

    menu_options = [
        { "option_name": "Story Mode", "on_select": story_mode },
        { "option_name": "Endless Mode", "on_select": endless_mode },
        { "option_name": "Save & Exit", "on_select": save_and_exit }
    ]

    return Menu(menu_options, game_state_parent)

def new_pause_menu(game_state_parent):
    def start_over(option, menu, game_state):
        menu.force_screen('puzzle')
        game_state.start_over()

    def hint(option, menu, game_state):
        menu.force_screen('puzzle')
        game_state.show_hint()

    def give_up(option, menu, game_state):
        menu.force_screen('puzzle')
        game_state.open_game_over_dialog()

    def how_to_play(option, menu, game_state):
        game_state.open_how_to_play_dialog()

    def settings(option, menu, game_state):
        settings_menu_options = [
            { "option_name": "Visual Themes", "on_select": visual_themes },
            { "option_name": "Quote File", "on_select": quote_file },
            { "option_name": "Sound Settings" }
        ]
        menu.open_sub_menu(settings_menu_options)

    def about(option, menu, game_state):
        game_state.open_about_dialog()

    def set_visual_theme(option, menu, game_state):
        game_state.user_preferences.visual_theme = option["option_name"]

    def set_quote_db(option, menu, game_state):
        game_state.user_preferences.quote_db = option["option_name"]
        menu.cancel()

    def visual_themes(option, menu, game_state):
        current_theme = game_state.user_preferences.visual_theme
        theme_list = [theme["id"] for theme in game_state.data_manager.read_theme_file()]
        theme_options = [{ "option_name": theme, "on_select": set_visual_theme } for theme in theme_list]
        current_theme_index = theme_list.index(current_theme) if current_theme in theme_list else 0
        menu.open_sub_menu(theme_options, cursor_index = current_theme_index)

    def quote_file(option, menu, game_state):
        current_quote_db = game_state.user_preferences.quote_db
        quote_db_list = [quote_db_name for quote_db_name in game_state.data_manager.list_quote_dbs()]
        quote_db_options = [{ "option_name": quote_db_name, "on_select": set_quote_db } for quote_db_name in quote_db_list]
        current_quote_db_index = quote_db_list.index(current_quote_db) if current_quote_db in quote_db_list else 0
        menu.open_sub_menu(quote_db_options, cursor_index = current_quote_db_index)

    if game_state_parent.mode == 'story':
        menu_options = [
            { "option_name": "Start Over", "on_select": start_over },
            { "option_name": "How To Play", "on_select": how_to_play },
            { "option_name": "Settings", "on_select": settings },
            { "option_name": "About", "on_select": about },
            { "option_name": "Quit to Main Menu", "on_select": quit_to_main_menu },
            { "option_name": "Save & Exit", "on_select": save_and_exit }
        ]
    else:
        menu_options = [
            { "option_name": "Start Over", "on_select": start_over },
            { "option_name": "Hint", "on_select": hint },
            { "option_name": "Give Up", "on_select": give_up },
            { "option_name": "How To Play", "on_select": how_to_play },
            { "option_name": "Settings", "on_select": settings },
            { "option_name": "About", "on_select": about },
            { "option_name": "Quit to Main Menu", "on_select": quit_to_main_menu },
            { "option_name": "Save & Exit", "on_select": save_and_exit }
        ]

    parent_screen = game_state_parent.screen

    return Menu(menu_options, game_state_parent, parent_screen = parent_screen)

def new_story_branch_menu(game_state_parent, choices):
    def select_story_branch(option, menu, game_state):
        game_state.load_story_level(game_state.story, option["option_value"], False)

    menu_options = [{ "option_name": i["name"], "on_select": select_story_branch, "option_value": i["next_level"] } for i in choices]

    return Menu(menu_options, game_state_parent)
